/**
 * Copyright (C) 2021 - 2031 O-Cubes Co., Ltd.
 */

/****************************************************************
 *  @file    sec.c
 *  @brief   Security Engine source file
 *  @version v1.0
 *  @date    03. Apr. 2023
 ****************************************************************/
#include <string.h>

#include "sec.h"
#include "printf.h"

/* AES/SM4 block size, 128bits */
#define AES_SM4_BLOCK_SIZE              (16)

/* software key length, 128bits */
#define SEC_KEY_BYTES_LEN               (16)
#define SEC_SW_KEY_UPDATE(j, sw_key)    (SEC_REG_PTR()->sw_key[j] = sw_key)

static inline void sec_en_aes_encrypt(void)
{
	SEC_REG_PTR()->cmd = AES_ENCRYPT;
}

static inline void sec_dis_aes_encrypt(void)
{
	SEC_REG_PTR()->cmd &= ~AES_ENCRYPT;
}

static inline void sec_en_aes_decrypt(void)
{
	SEC_REG_PTR()->cmd = AES_DECRYPT;
}

static inline void sec_dis_aes_decrypt(void)
{
	SEC_REG_PTR()->cmd &= ~AES_DECRYPT;
}

static inline void sec_en_sm4_encrypt(void)
{
	SEC_REG_PTR()->cmd = SM4_ENCRYPT;
}

static inline void sec_dis_sm4_encrypt(void)
{
	SEC_REG_PTR()->cmd &= ~SM4_ENCRYPT;
}

static inline void sec_en_sm4_decrypt(void)
{
	SEC_REG_PTR()->cmd = SM4_DECRYPT;
}

static inline void sec_dis_sm4_decrypt(void)
{
	SEC_REG_PTR()->cmd &= ~SM4_DECRYPT;
}

static inline void sec_en_sha256(void)
{
	SEC_REG_PTR()->cmd = SHA256;
}

static inline void sec_dis_sha256(void)
{
	SEC_REG_PTR()->cmd &= ~SHA256;
}

void sec_update_hw_key(void)
{
	u32 val = SEC_REG_PTR()->cmd;

	val |= UPDATE_KEY;
	SEC_REG_PTR()->cmd = val;

	val &= ~UPDATE_KEY;
	SEC_REG_PTR()->cmd = val;
}

void sec_set_key_type(enum key_type type)
{
	u32 val = SEC_REG_PTR()->cfg;

	val = SET_REG_BITS(val, KEY_TYPE_POS, KEY_TYPE_LEN, type);

	SEC_REG_PTR()->cfg = val;
}

static inline void sec_set_cmd_data_len(u32 len)
{
	SEC_REG_PTR()->cmd_size = len;
}

static inline s32 sec_wait_cfg_done(u32 timeout)
{
	return wait_cond_timeout(SEC_REG_PTR()->mis & SEC_CFG, timeout);
}

static inline void sec_clr_cfg_status(void)
{
	SEC_REG_PTR()->icr |= SEC_CFG;
}

static inline s32 sec_wait_done_status(u32 timeout)
{
	return wait_cond_timeout(SEC_REG_PTR()->mis & SEC_DONE, timeout);
}

static inline void sec_clr_done_status(void)
{
	SEC_REG_PTR()->icr |= SEC_DONE;
	SEC_REG_PTR()->icr &= ~SEC_DONE;
}

static int sec_update_sw_key(const u8 *key, int len)
{
	int i;
	int j;
	u32 sw_key;

	if (len != SEC_KEY_BYTES_LEN)
		return -1;

	for (i = 0, j = 0; i < SEC_KEY_BYTES_LEN; i += 4) {
		sw_key = key[i] | key[i + 1] << 8 | key[i + 2] << 16 | key[i + 3] << 24;
		SEC_SW_KEY_UPDATE(j++, sw_key);
	}

	return 0;
}

void sec_set_rng_seed(int seed)
{
	u32 val = SEC_REG_PTR()->rng_seed;

	val = SET_REG_BITS(val, RNG_SEED_POS, RNG_SEED_LEN, seed);

	SEC_REG_PTR()->rng_seed = val;
}

void sec_en_rng(void)
{
	SEC_REG_PTR()->rng |= RNG_SWITCH_ON;
}

void sec_dis_rng(void)
{
	SEC_REG_PTR()->rng &= ~RNG_SWITCH_ON;
}

void sec_update_rng_seed(void)
{
	SEC_REG_PTR()->rng |= RNG_UPDATE_SEED;
}

void sec_get_hw_key(u8 *hw_key, int keylen)
{
	memcpy(hw_key, (void *)SEC_HW_KEY_DATA_BASE_ADDR, keylen);
}

void sec_get_rng_data(u8 *rng, int rng_len)
{
	memcpy(rng, (void *)SEC_RANDOM_DATA_BASE_ADDR, rng_len);
}

inline u32 sec_get_proc_len(void)
{
	return SEC_REG_PTR()->proc_len;
}

int sec_aes_set_ek(const u8 *ek, int keylen)
{
	sec_set_key_type(SW_KEY);

	return sec_update_sw_key(ek, keylen);
}

int sec_aes_set_dk(const u8 *dk, int keylen)
{
	sec_set_key_type(SW_KEY);

	return sec_update_sw_key(dk, keylen);
}

int sec_aes_ecb_encrypt(const u8 *input, int length, u8 *output)
{
	int i;
	int len;
	int rest = length;

	for (i = 0; i < length;) {
		len = (rest > AES_SM4_BLOCK_SIZE) ? AES_SM4_BLOCK_SIZE : rest;

		sec_set_cmd_data_len(len);
		sec_en_aes_encrypt();

		if (sec_wait_cfg_done(10))
			return -1;

		sec_clr_cfg_status();

		memcpy((void *)SEC_OUT_DATA_BASE_ADDR, input + i, len);

		if (sec_wait_done_status(10))
			return -1;

		sec_clr_done_status();

		memcpy(output + i, (void *)SEC_OUT_DATA_BASE_ADDR, len);
		i += len;
		rest -= len;

		sec_dis_aes_encrypt();
	}

	return 0;
}

int sec_aes_ecb_decrypt(const u8 *input, int length, u8 *output)
{
	int i;
	int len;
	int rest = length;

	for (i = 0; i < length;) {
		len = (rest > AES_SM4_BLOCK_SIZE) ? AES_SM4_BLOCK_SIZE : rest;

		sec_set_cmd_data_len(len);
		sec_en_aes_decrypt();

		if (sec_wait_cfg_done(10))
			return -1;

		sec_clr_cfg_status();

		memcpy((void *)SEC_OUT_DATA_BASE_ADDR, input + i, len);

		if (sec_wait_done_status(10))
			return -1;

		sec_clr_done_status();

		memcpy(output + i, (void *)SEC_OUT_DATA_BASE_ADDR, len);
		i += len;
		rest -= len;

		sec_dis_aes_decrypt();
	}

	return 0;
}

int sec_sm4_set_ek(const u8 *ek, int keylen)
{
	sec_set_key_type(SW_KEY);

	return sec_update_sw_key(ek, keylen);
}

int sec_sm4_set_dk(const u8 *dk, int keylen)
{
	sec_set_key_type(SW_KEY);

	return sec_update_sw_key(dk, keylen);
}

int sec_sm4_ecb_encrypt(const u8 *input, int length, u8 *output)
{
	int i;
	int len;
	int rest = length;

	for (i = 0; i < length;) {
		len = (rest > AES_SM4_BLOCK_SIZE) ? AES_SM4_BLOCK_SIZE : rest;

		sec_set_cmd_data_len(len);
		sec_en_sm4_encrypt();

		if (sec_wait_cfg_done(10))
			return -1;

		sec_clr_cfg_status();

		memcpy((void *)SEC_OUT_DATA_BASE_ADDR, input + i, len);

		if (sec_wait_done_status(10))
			return -1;

		sec_clr_done_status();

		memcpy(output + i, (void *)SEC_OUT_DATA_BASE_ADDR, len);
		i += len;
		rest -= len;

		sec_dis_sm4_encrypt();
	}

	return 0;
}

int sec_sm4_ecb_decrypt(const u8 *input, int length, u8 *output)
{
	int i;
	int len;
	int rest = length;

	for (i = 0; i < length;) {
		len = (rest > AES_SM4_BLOCK_SIZE) ? AES_SM4_BLOCK_SIZE : rest;

		sec_set_cmd_data_len(len);
		sec_en_sm4_decrypt();

		if (sec_wait_cfg_done(10))
			return -1;

		sec_clr_cfg_status();

		memcpy((void *)SEC_OUT_DATA_BASE_ADDR, input + i, len);

		if (sec_wait_done_status(10))
			return -1;

		sec_clr_done_status();

		memcpy(output + i, (void *)SEC_OUT_DATA_BASE_ADDR, len);
		i += len;
		rest -= len;

		sec_dis_sm4_decrypt();
	}

	return 0;
}

int sec_sha256_compute(const u8 *data, u32 length, u8 *digest)
{
	int i;
	int len;
	int rest = length;

	if (length > SHA256_MAXIMUM_LEN)
		return -1;

	sec_set_cmd_data_len(length);

	sec_en_sha256();

	if (sec_wait_cfg_done(10))
		return -1;

	sec_clr_cfg_status();

	for (i = 0; i < length;) {
		len = (rest > SEC_MAX_DATA_LEN) ? SEC_MAX_DATA_LEN : rest;

		memcpy((void *)SEC_IN_DATA_BASE_ADDR, data + i, len);
		i += len;
		rest -= len;
	}

	if (sec_wait_done_status(10))
		return -1;

	sec_clr_done_status();

	memcpy(digest, (void *)SEC_OUT_DATA_BASE_ADDR, 256 / 8);

	sec_dis_sha256();

	return 0;
}

